import torch
import random
import numpy as np


from PIL import Image
from os.path import join
import torch.utils.data as data
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# this is default dataset dir, this path will be change by the main program
structfile = '/home/a409/users/lihaowei/data/CVUSA_DSM/'


class LimitedFoV(object):
    def __init__(self, fov=360.):
        self.fov = fov

    def __call__(self, x):
        angle = random.randint(0, 359)
        rotate_index = int(angle / 360. * x.shape[2])
        fov_index = int(self.fov / 360. * x.shape[2])
        if rotate_index > 0:
            img_shift = torch.zeros(x.shape)
            img_shift[:, :, :rotate_index] = x[:, :, -rotate_index:]
            img_shift[:, :, rotate_index:] = x[:, :, :(x.shape[2] - rotate_index)]
        else:
            img_shift = x
        return img_shift[:, :, :fov_index], angle


class LimitedFoV2(object):
    def __init__(self, fov=360.):
        self.fov = fov

    def __call__(self, x, angle):
        rotate_index = int(angle / 360. * x.shape[2])
        fov_index = int(self.fov / 360. * x.shape[2])
        if rotate_index > 0:
            img_shift = torch.zeros(x.shape)
            img_shift[:, :, :rotate_index] = x[:, :, -rotate_index:]
            img_shift[:, :, rotate_index:] = x[:, :, :(x.shape[2] - rotate_index)]
        else:
            img_shift = x
        return img_shift[:, :, :fov_index]


class Expend(object):
    def __init__(self, fov=360.):
        self.fov = fov
    def __call__(self, x):
        fov_index = int(self.fov / 360. * x.shape[2])
        img_shift = x[:, :, :- fov_index]
        return torch.cat((x, img_shift), -1)

def input_transform():
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
    ])


def get_whole_training_set(mode, angle):
    return WholeDatasetFromStruct(structfile, mode, angle)


def get_train2_set(mode, angle):
    return DataForStep2(structfile,  mode, angle)


def get_whole_test_set(mode, angle):
    return WholeDatasetFromStruct(structfile,  mode, angle)


class WholeDatasetFromStruct(data.Dataset):
    def __init__(self, structfile, mode, angle):
        super().__init__()

        self.img_root = structfile
        self.data_list = self.img_root + 'splits/'+str(mode)+'.csv'
        self.input_transform = input_transform()
        self.fov = LimitedFoV(angle)

        self.__cur_id = 0
        self.sat_list = []
        self.grd_list = []
        self.id_list = []
        self.id_idx_list = []
        with open(self.data_list, 'r') as file:
            idx = 0
            for line in file:
                line = line.replace('\"', '')
                data = line.split(',')
                pano_id = (data[0].split('/')[-1]).split('.')[0]
                # satellite filename, streetview filename, pano_id
                item = data[0].replace('bing', 'polar').replace('jpg', 'png')
                grd_inx = data[1]
                self.sat_list.append(item)
                self.grd_list.append(grd_inx)
                self.id_list.append(pano_id)
                self.id_idx_list.append(idx)
                idx += 1
        self.data_size = len(self.id_list)

    def __getitem__(self, index):
        sat_img = Image.open(join(structfile, self.sat_list[index]))
        grd_img = Image.open(join(structfile, self.grd_list[index]))

        sat_img = self.input_transform(sat_img)
        grd_img = self.input_transform(grd_img)
        grd_img, angle = self.fov(grd_img)
        return sat_img, grd_img, self.sat_list[index], str(angle)

    def __len__(self):  # 返回列表长度
        return len(self.sat_list)


class DataForStep2(data.Dataset):
    def __init__(self, structfile, mode, fov):
        super().__init__()
        self.mode = mode
        self.img_root = structfile
        self.input_transform = input_transform()
        self.fov = LimitedFoV2(fov)
        self.expend = Expend(fov)

        if self.mode == 'train':
            self.part2_angle_ = self.img_root + 'part2_splits/part2_angle.txt'
            self.part2_index_ = self.img_root + 'part2_splits/part2_index.txt'
            self.part2_index_sat_ = self.img_root + 'part2_splits/part2_index_sat.txt'
            self.part2_path_ = self.img_root + 'part2_splits/part2_path.txt'
        else:
            self.part2_angle_ = self.img_root + 'part2_splits/cache/part2_angle.txt'
            self.part2_index_ = self.img_root + 'part2_splits/cache/part2_index.txt'
            self.part2_index_sat_ = self.img_root + 'part2_splits/cache/part2_index_sat.txt'
            self.part2_path_ = self.img_root + 'part2_splits/cache/part2_path.txt'

        f = open(self.part2_angle_, 'r')
        self.part2_angle = f.readlines()
        f = open(self.part2_path_, 'r')
        self.part2_path = f.readlines()
        f = open(self.part2_index_, 'r')
        self.part2_index = f.readlines()
        f = open(self.part2_index_sat_, 'r')
        self.part2_index_sat = f.readlines()

        self.part2_angle = list(map(lambda x: int(x[1:-1]), self.part2_angle[0][1:-2].split(', ')))
        self.part2_path = list(map(lambda x: x[13:-5], self.part2_path[0][1:-2].split(', ')))
        self.part2_index = list(map(lambda x: int(x[:-1]), self.part2_index))
        self.part2_index_sat = list(
            map(lambda x: list(map(lambda y: int(y), x[1:-2].split(', '))), self.part2_index_sat))

        self.angle = []
        self.grd_list = []
        self.sat_list = []
        self.correct_location = []

        for i, index in enumerate(self.part2_index):
            self.correct_location.append(self.part2_index_sat[i].index(self.part2_index[i]))
            self.angle.append(self.part2_angle[index])
            self.grd_list.append('streetview/panos/' + self.part2_path[index] + '.jpg')
            self.sat_list.append(
                list(map(lambda x: 'polarmap/19/' + self.part2_path[x] + '.png', self.part2_index_sat[i])))

        self.size = len(self.grd_list)

        del self.part2_angle, self.part2_path
        del self.part2_index, self.part2_index_sat

    def __getitem__(self, index):
        sat_img = []
        grd_img = Image.open(join(structfile, self.grd_list[index]))

        # Fix me : use torch.tensor to rewrite
        for i in self.sat_list[index]:
            temp = Image.open(join(structfile, i))
            temp = self.input_transform(temp)
            temp = self.expend(temp)
            sat_img.append(temp)
            # temp = torch.unsqueeze(temp, 0)
            # sat_img = torch.cat((sat_img, temp), 0)
        grd_img = self.input_transform(grd_img)
        grd_img = self.fov(grd_img, self.angle[index])
        return torch.stack(sat_img), grd_img, self.angle[index], self.correct_location[index]

    def __len__(self):  # 返回列表长度
        return self.size


# test
if __name__ == '__main__':
    a = torch.zeros((3,100,200))
    b = Expend(70)
    c = b(a)
    print(c.shape)




